{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Week 2 Day 3\n",
    "\n",
    "Now we get to more detail:\n",
    "\n",
    "1. Different models\n",
    "\n",
    "2. Structured Outputs\n",
    "\n",
    "3. Guardrails"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "from openai import AsyncOpenAI\n",
    "from agents import Agent, Runner, trace, function_tool, OpenAIChatCompletionsModel, input_guardrail, GuardrailFunctionOutput\n",
    "from typing import Dict\n",
    "import requests\n",
    "import os\n",
    "from pydantic import BaseModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv(override=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "openai_api_key = os.getenv('OPENAI_API_KEY')\n",
    "google_api_key = os.getenv('GOOGLE_API_KEY')\n",
    "deepseek_api_key = os.getenv('DEEPSEEK_API_KEY')\n",
    "groq_api_key = os.getenv('GROQ_API_KEY')\n",
    "\n",
    "if openai_api_key:\n",
    "    print(f\"OpenAI API Key exists and begins {openai_api_key[:8]}\")\n",
    "else:\n",
    "    print(\"OpenAI API Key not set\")\n",
    "\n",
    "if google_api_key:\n",
    "    print(f\"Google API Key exists and begins {google_api_key[:2]}\")\n",
    "else:\n",
    "    print(\"Google API Key not set (and this is optional)\")\n",
    "\n",
    "if deepseek_api_key:\n",
    "    print(f\"DeepSeek API Key exists and begins {deepseek_api_key[:3]}\")\n",
    "else:\n",
    "    print(\"DeepSeek API Key not set (and this is optional)\")\n",
    "\n",
    "if groq_api_key:\n",
    "    print(f\"Groq API Key exists and begins {groq_api_key[:4]}\")\n",
    "else:\n",
    "    print(\"Groq API Key not set (and this is optional)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "instructions1 = \"You are a sales agent working for ComplAI, \\\n",
    "a company that provides a SaaS tool for ensuring SOC2 compliance and preparing for audits, powered by AI. \\\n",
    "You write professional, serious cold emails.\"\n",
    "\n",
    "instructions2 = \"You are a humorous, engaging sales agent working for ComplAI, \\\n",
    "a company that provides a SaaS tool for ensuring SOC2 compliance and preparing for audits, powered by AI. \\\n",
    "You write witty, engaging cold emails that are likely to get a response.\"\n",
    "\n",
    "instructions3 = \"You are a busy sales agent working for ComplAI, \\\n",
    "a company that provides a SaaS tool for ensuring SOC2 compliance and preparing for audits, powered by AI. \\\n",
    "You write concise, to the point cold emails.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### It's easy to use any models with OpenAI compatible endpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "GEMINI_BASE_URL = \"https://generativelanguage.googleapis.com/v1beta/openai/\"\n",
    "DEEPSEEK_BASE_URL = \"https://api.deepseek.com/v1\"\n",
    "GROQ_BASE_URL = \"https://api.groq.com/openai/v1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "deepseek_client = AsyncOpenAI(base_url=DEEPSEEK_BASE_URL, api_key=deepseek_api_key)\n",
    "gemini_client = AsyncOpenAI(base_url=GEMINI_BASE_URL, api_key=google_api_key)\n",
    "groq_client = AsyncOpenAI(base_url=GROQ_BASE_URL, api_key=groq_api_key)\n",
    "\n",
    "deepseek_model = OpenAIChatCompletionsModel(model=\"deepseek-chat\", openai_client=deepseek_client)\n",
    "gemini_model = OpenAIChatCompletionsModel(model=\"gemini-2.0-flash\", openai_client=gemini_client)\n",
    "llama3_3_model = OpenAIChatCompletionsModel(model=\"llama-3.3-70b-versatile\", openai_client=groq_client)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use different models, falling back to OpenAI if others don't work\n",
    "sales_agent1 = Agent(name=\"OpenAI Sales Agent\", instructions=instructions1, model=\"gpt-4o-mini\")\n",
    "sales_agent2 = Agent(name=\"OpenAI Sales Agent 2\", instructions=instructions2, model=\"gpt-4o\")  \n",
    "sales_agent3 = Agent(name=\"OpenAI Sales Agent 3\", instructions=instructions3, model=\"gpt-3.5-turbo\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "description = \"Write a cold sales email\"\n",
    "\n",
    "tool1 = sales_agent1.as_tool(tool_name=\"sales_agent1\", tool_description=description)\n",
    "tool2 = sales_agent2.as_tool(tool_name=\"sales_agent2\", tool_description=description)\n",
    "tool3 = sales_agent3.as_tool(tool_name=\"sales_agent3\", tool_description=description)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@function_tool\n",
    "def send_simple_notification(message: str) -> Dict[str, str]:\n",
    "    \"\"\" Send a simple notification with the given message to all sales prospects \"\"\"\n",
    "    url = \"https://api.pushover.net/1/messages.json\"\n",
    "    \n",
    "    data = {\n",
    "        \"token\": os.environ.get('PUSHOVER_TOKEN'),\n",
    "        \"user\": os.environ.get('PUSHOVER_USER'),\n",
    "        \"message\": message,\n",
    "        \"title\": \"Sales Email\"\n",
    "    }\n",
    "    \n",
    "    response = requests.post(url, data=data)\n",
    "    if response.status_code == 200:\n",
    "        return {\"status\": \"success\", \"message\": \"Notification sent successfully\"}\n",
    "    else:\n",
    "        return {\"status\": \"error\", \"message\": f\"Error {response.status_code}: {response.text}\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subject_instructions = \"You can write a subject for a cold sales email. \\\n",
    "You are given a message and you need to write a subject for an email that is likely to get a response.\"\n",
    "\n",
    "html_instructions = \"You can convert a text email body to an HTML email body. \\\n",
    "You are given a text email body which might have some markdown \\\n",
    "and you need to convert it to an HTML email body with simple, clear, compelling layout and design.\"\n",
    "\n",
    "subject_writer = Agent(name=\"Email subject writer\", instructions=subject_instructions, model=\"gpt-4o-mini\")\n",
    "subject_tool = subject_writer.as_tool(tool_name=\"subject_writer\", tool_description=\"Write a subject for a cold sales email\")\n",
    "\n",
    "html_converter = Agent(name=\"HTML email body converter\", instructions=html_instructions, model=\"gpt-4o-mini\")\n",
    "html_tool = html_converter.as_tool(tool_name=\"html_converter\",tool_description=\"Convert a text email body to an HTML email body\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "notification_tools = [send_simple_notification]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "instructions = \"You are a notification sender. You receive the body of an email to be sent. \\\n",
    "You use the send_simple_notification tool to send the email content as a clean, readable notification. \\\n",
    "Do not add any formatting or subjects - just send the email content directly.\"\n",
    "\n",
    "\n",
    "notification_agent = Agent(\n",
    "    name=\"Notification Sender\",\n",
    "    instructions=instructions,\n",
    "    tools=[send_simple_notification],\n",
    "    model=\"gpt-4o-mini\",\n",
    "    handoff_description=\"Send the email content as a simple notification\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = [tool1, tool2, tool3]\n",
    "handoffs = [notification_agent]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sales_manager_instructions = \"\"\"\n",
    "You are a Sales Manager at ComplAI. Your goal is to find the single best cold sales email using the sales_agent tools.\n",
    "\n",
    "CRITICAL: You must complete ALL steps before making ANY handoff. You can only make ONE handoff total.\n",
    "\n",
    "Follow these steps carefully:\n",
    "1. Generate Drafts: Use all three sales_agent tools to generate three different email drafts. Do not proceed until all three drafts are ready.\n",
    "\n",
    "2. Evaluate and Select: Review ALL three drafts and choose the single best email using your judgment of which one is most effective. DO NOT hand off during this step.\n",
    "\n",
    "3. Single Handoff: Hand off ONLY the winning email draft text to the 'Notification Sender' agent. \n",
    "\n",
    "FORBIDDEN ACTIONS:\n",
    "- DO NOT hand off each email as you generate it\n",
    "- DO NOT make multiple handoffs\n",
    "- DO NOT hand off all three emails\n",
    "- You can make ONLY ONE handoff with ONLY the best email text\n",
    "\n",
    "Crucial Rules:\n",
    "- You must use the sales agent tools to generate the drafts — do not write them yourself.\n",
    "- You must hand off exactly ONE email to the Notification Sender — never more than one.\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "sales_manager = Agent(\n",
    "    name=\"Sales Manager\",\n",
    "    instructions=sales_manager_instructions,\n",
    "    tools=tools,\n",
    "    handoffs=[notification_agent],\n",
    "    model=\"gpt-4o-mini\")\n",
    "\n",
    "message = \"Send out a cold sales email addressed to Dear CEO from Alice\"\n",
    "\n",
    "with trace(\"Automated SDR\"):\n",
    "    result = await Runner.run(sales_manager, message, max_turns=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check out the trace:\n",
    "\n",
    "https://platform.openai.com/traces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NameCheckOutput(BaseModel):\n",
    "    is_name_in_message: bool\n",
    "    name: str\n",
    "\n",
    "guardrail_agent = Agent( \n",
    "    name=\"Name check\",\n",
    "    instructions=\"Check if the user is including someone's personal name in what they want you to do.\",\n",
    "    output_type=NameCheckOutput,\n",
    "    model=\"gpt-4o-mini\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@input_guardrail\n",
    "async def guardrail_against_name(ctx, agent, message):\n",
    "    result = await Runner.run(guardrail_agent, message, context=ctx.context)\n",
    "    is_name_in_message = result.final_output.is_name_in_message\n",
    "    return GuardrailFunctionOutput(output_info={\"found_name\": result.final_output},tripwire_triggered=is_name_in_message)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "careful_sales_manager = Agent(\n",
    "    name=\"Careful Sales Manager\",\n",
    "    instructions=sales_manager_instructions,\n",
    "    tools=tools,\n",
    "    handoffs=[notification_agent],\n",
    "    model=\"gpt-4o-mini\",\n",
    "    input_guardrails=[guardrail_against_name]\n",
    "    )\n",
    "\n",
    "# Test Guardrail Behavior - Choose one message:\n",
    "\n",
    "# Option 1: Message WITHOUT personal name - Guardrail will allow execution\n",
    "# message = \"Send out a cold sales email addressed to Dear CEO from Sales Team\"\n",
    "\n",
    "# Option 2: Message WITH personal name - Guardrail will BLOCK with InputGuardrailTripwireTriggered\n",
    "message = \"Send out a cold sales email addressed to Dear CEO from Alice\"\n",
    "\n",
    "print(f\"Testing with message: {message}\")\n",
    "if \"Alice\" in message or \"from \" in message.split(\"from \")[-1] and len(message.split(\"from \")[-1].split()) == 1:\n",
    "    print(\"Expected: Guardrail will BLOCK execution (personal name detected)\")\n",
    "else:\n",
    "    print(\"Expected: Guardrail allows execution (no personal name detected)\")\n",
    "\n",
    "with trace(\"Protected Automated SDR\"):\n",
    "    result = await Runner.run(careful_sales_manager, message, max_turns=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check out the trace:\n",
    "\n",
    "https://platform.openai.com/traces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# This message should pass the guardrail (no personal name)\n",
    "message = \"Send out a cold sales email addressed to Dear CEO from Head of Business Development\"\n",
    "\n",
    "print(f\"Testing with message: {message}\")\n",
    "print(\"Expected: Guardrail allows execution (no personal name detected)\")\n",
    "\n",
    "with trace(\"Protected Automated SDR\"):\n",
    "    result = await Runner.run(careful_sales_manager, message, max_turns=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<table style=\"margin: 0; text-align: left; width:100%\">\n",
    "    <tr>\n",
    "        <td style=\"width: 150px; height: 150px; vertical-align: middle;\">\n",
    "            <img src=\"../assets/exercise.png\" width=\"150\" height=\"150\" style=\"display: block;\" />\n",
    "        </td>\n",
    "        <td>\n",
    "            <h2 style=\"color:#ff7800;\">Exercise</h2>\n",
    "            <span style=\"color:#ff7800;\">• Try different models<br/>• Add more input and output guardrails<br/>• Use structured outputs for the email generation\n",
    "            </span>\n",
    "        </td>\n",
    "    </tr>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
